package com.ruoyi.ade.llm.service.impl;

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.ade.llm.agent.MilvusAgent;
import com.ruoyi.ade.llm.agent.OllamaAgent;
import com.ruoyi.ade.llm.agent.QianfanAgent;
import com.ruoyi.ade.llm.chunk.*;
import com.ruoyi.ade.llm.exception.AdeKnowledgeException;
import com.ruoyi.ade.llm.model.*;
import com.ruoyi.ade.llm.service.IAdeKnowledgeEmbedService;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
@Slf4j
public class AdeKnowledgeEmbedServiceImpl implements IAdeKnowledgeEmbedService {

    //ade的向量知识库 集合
//    private final String collection = "adesv2_pos";
//    private final String model_embedding = EmbeddingRequest.MODEL_DMETS;  //采用BGE向量模型
//    private final String model_rerank = "bce-reranker-base";
    //需要使用到Ollama的服务  主要是做向量化处理
    private OllamaAgent ollamaAgent = OllamaAgent.getInstance();
    //需要用到Milvus服务
    private MilvusAgent milvusAgent = MilvusAgent.getInstance();
    //需要用到Qianfan服务 主要做rerank
    private QianfanAgent qianfanAgent = QianfanAgent.getInstance();

    public void createCollection(String collection) {
        milvusAgent.createCollection(collection);
    }

    public boolean dropCollection(String collection){
         return milvusAgent.dropCollection(collection);
    }

    public void index(String collection, String embedModel, AdeKnowledge ade) {
        //Step 1: convert embedding by ollma
        EmbeddingRequest embeddingRequest = new EmbeddingRequest(embedModel, Arrays.asList(new String[]{ade.getText()}));  //仅仅请求单个text向量化
        EmbeddingResponse embeddingResponse = ollamaAgent.embed(embeddingRequest);
        Float[] vector = embeddingResponse.getEmbeddings()[0]; //没有采用批量想量化，仅仅提供了一条text  所以返回的向量取第一个

        //Step2: store to milvus
        AdeMilvusRequest obj = new AdeMilvusRequest(ade.getId(), vector, ade.getText(), ade.getAde(), ade.getExplain());
        milvusAgent.index(collection, obj);
        System.out.println("索引成功：" + ade.getText());
        log.info("Collection-{} index ade:{} successfully! ", collection, ade.getId());
    }

    public void index(String collection, String embedModel, List<AdeKnowledge> ades) {
        if (ades != null) {
            for (AdeKnowledge ade : ades) {
                index(collection, embedModel, ade);
            }
        }
    }

    @Override
    public void index(String collection, String embedModel, List<AdeKnowledge> ades, float simScoreThreshold) {
        if (ades != null) {
            for (AdeKnowledge ade : ades) {
                index(collection, embedModel, ade, simScoreThreshold);
            }
        }
    }

    @Override
    public void index(String collection, String embedModel, AdeKnowledge ade, float simScoreThreshold) {
        //Step 1: convert embedding by ollma
        EmbeddingRequest embeddingRequest = new EmbeddingRequest(embedModel, Arrays.asList(new String[]{ade.getText()}));  //仅仅请求单个text向量化
        EmbeddingResponse embeddingResponse = ollamaAgent.embed(embeddingRequest);
        Float[] vector = embeddingResponse.getEmbeddings()[0]; //没有采用批量想量化，仅仅提供了一条text  所以返回的向量取第一个

        //Step2: 判断是否符合插入条件
        List<AdeMilvusResponse> milvusObjectResponses = milvusAgent.search(collection, vector, simScoreThreshold, 5);  //找到Top20 满足阈值的知识
        if (milvusObjectResponses.size() > 0) {
            List<String> docs = new ArrayList<>();
            for (AdeMilvusResponse resp : milvusObjectResponses) {
                docs.add(resp.getText());
            }
//            RerankRequest rerankRequest = new RerankRequest(model_rerank, ade.getText(), docs, docs.size());
//            RerankResponse rerankResponse = qianfanAgent.reran(rerankRequest);
//                    for (RerankResponse.RerankDocument doc : rerankResponse.getResults()) {
//                        System.out.println(doc.getRelevance_score()+"\t\t"+doc.getDocument());
//                    }

            throw new AdeKnowledgeException("存在重复/相似的知识!");
        }
//                if(milvusObjectResponses.size()>0) {  //找到相似性的知识了   则需要进一步计算它们之间的相关性
//                    List<RetrivedAdeKnowledge> raks = new ArrayList<>();
//                    Map<String, RetrivedAdeKnowledge> map = new HashMap<>();
//                    for (MilvusObjectResponse resp : milvusObjectResponses) {
//                        RetrivedAdeKnowledge r = RetrivedAdeKnowledge.fromMilvusObject(resp);
//                        map.put(r.getKnowledge().getText(), r);
//                    }
//                    List<String> docs = new ArrayList<>();
//                    for (MilvusObjectResponse resp : milvusObjectResponses) docs.add(resp.getText());
//                    RerankRequest rerankRequest = new RerankRequest(model_rerank, ade.getText(), docs, 20);
//                    RerankResponse rerankResponse = qianfanAgent.reran(rerankRequest);
//                    for (RerankResponse.RerankDocument doc : rerankResponse.getResults()) {
//                        RetrivedAdeKnowledge tmp = map.get(doc.getDocument());
//                        tmp.setReScore(doc.getRelevance_score());
//                    }
//                    for(RetrivedAdeKnowledge rak: map.values()){
//                        if(rak.getSimScore()>=config.getSimScoreThreshold() &&  rak.getReScore()>=config.getReScoreThreshold()){
//                            //如果相似性得分 以及相关性得分都符合要求
//                        }
//                    }
//
//                }
        //Step2: store to milvus
        AdeMilvusRequest obj = new AdeMilvusRequest(ade.getId(), vector, ade.getText(), ade.getAde(), ade.getExplain());
        milvusAgent.index(collection, obj);
//        System.out.println("索引成功：" + ade.getText());
    }

    @Override
    public AdeRetrivedResponse search(String collection, String embedModel, String emr, String chunkType, float score, int topN, boolean rerank, boolean debug) {

        AdeRetrivedResponse response = new AdeRetrivedResponse();
        response.setPairs(new ArrayList<>());
        response.setAdes(new ArrayList<>());


        List<String> references = new ArrayList<>();
        IAbstractChunkRequest chunkRequest = null;
        //Step 1: 判断是否要做分割
        if (IAbstractChunkRequest.SPLIT_PATTERN.equals(chunkType)) {
            chunkRequest = new ChunkByPatternRequest();
            IChunkService chunkService = new ChunkByPatternService();
            List<String> sentences = chunkService.chunk(emr, chunkRequest);
            references.addAll(sentences);
        } else if (IAbstractChunkRequest.SPLIT_MARK.equals(chunkType)) {
            chunkRequest = ChunkByMarkRequest.getDefault();
            IChunkService chunkService = new ChunkByMarkService();
            List<String> sentences = chunkService.chunk(emr, chunkRequest);
            references.addAll(sentences);
        } else if (IAbstractChunkRequest.SPLIT_NONE.equals(chunkType)) {
            references.add(emr);
        } else {
            throw new AdeKnowledgeException("Unknown chunk type:" + chunkType);
        }

        //step2: 每个chunk分析
        if(debug) {   //如果是debug模式则要要分析每个拆分的句子，同时找到每个拆分句子对应的结果
            for (String reference : references) {
                long cur = System.currentTimeMillis();
                AdeRetrivedPair pair = new AdeRetrivedPair();
                pair.setChunk(reference);
                pair.setAdes(new ArrayList<>());
                EmbeddingRequest embeddingRequest = new EmbeddingRequest(embedModel, Arrays.asList(new String[]{reference}));
                EmbeddingResponse embeddingResponse = ollamaAgent.embed(embeddingRequest);
                Float[][] vector = embeddingResponse.getEmbeddings();
                System.out.println("向量化耗时:" + (System.currentTimeMillis() - cur));
                cur = System.currentTimeMillis();
                //从向量库批量检索
                List<AdeMilvusResponse> responses = milvusAgent.search(collection, vector, score, topN);
                System.out.println(reference + "向量索引....耗时:\t\t" + (System.currentTimeMillis() - cur));
                for (AdeMilvusResponse resp : responses) {
                    AdeRetrivedKnowledge r = AdeRetrivedKnowledge.fromMilvusObject(resp);
                    pair.getAdes().add(r);
                    response.getAdes().add(r);
                }
                Collections.sort(pair.getAdes()); //排序 从大到小按分数排序
                response.getPairs().add(pair);
            }
        }else{
            EmbeddingRequest embeddingRequest1 = new EmbeddingRequest(embedModel, references);
            EmbeddingResponse embeddingResponse1 = ollamaAgent.embed(embeddingRequest1);
            Float[][] vector1 = embeddingResponse1.getEmbeddings();
            List<AdeMilvusResponse> responses1 = milvusAgent.search(collection, vector1, score, topN);
            for (AdeMilvusResponse resp : responses1) {
                AdeRetrivedKnowledge r = AdeRetrivedKnowledge.fromMilvusObject(resp);
                response.getAdes().add(r);
            }
        }

//        //Step2: 基于向量数据库 检索相似性ADE知识
//        //将查询的文本转换成向量
//        EmbeddingRequest embeddingRequest = new EmbeddingRequest(embedModel, references);  //批量请求向量化
//        EmbeddingResponse embeddingResponse = ollamaAgent.embed(embeddingRequest);
//        Float[][] vector = embeddingResponse.getEmbeddings();
//        //从向量库批量检索
//        List<AdeMilvusResponse> responses = milvusAgent.search(collection, vector, score, topN);
//        if (responses.size() > topN) responses = responses.subList(0, topN);
//
//        //Step4: 组装RetrievedAcknowledge回去
//        List<AdeRetrivedKnowledge> raks = new ArrayList<>();
//        Map<String, AdeRetrivedKnowledge> map = new HashMap<>();
//        for (AdeMilvusResponse resp : responses) {
//            AdeRetrivedKnowledge r = AdeRetrivedKnowledge.fromMilvusObject(resp);
//            map.put(r.getKnowledge().getText(), r);
//        }

        //Step3: 判断是否需要Rerank
//        if (rerank) {
//            List<String> docs = new ArrayList<>();
//            for (MilvusObjectResponse resp : responses) docs.add(resp.getText());
//            RerankRequest rerankRequest = new RerankRequest(model_rerank, emr, docs, topN);
//            RerankResponse rerankResponse = qianfanAgent.reran(rerankRequest);
//            for (RerankResponse.RerankDocument doc : rerankResponse.getResults()) {
//                AdeRetrivedKnowledge tmp = map.get(doc.getDocument());
//                tmp.setReScore(doc.getRelevance_score());
//                raks.add(tmp);
//            }
//        } else {
//            raks.addAll(map.values());
//            Collections.sort(raks);
//        }
//            Collections.sort(response.getAdes());  //排序 从大到小按分数排序
        response.setAdes(filterRepartAndSort(response.getAdes(),topN));
        return response;
    }

    @Override
    public AdeRetrivedResponse search(String emr, AdeKnowledgeRetriveConfig retriveConfig, AdeKnowledgeRerankConfig rerankConfig) {

        if (StringUtils.isEmpty(emr)) throw new AdeKnowledgeException("待分析的病历文书不能为空!");
        if (retriveConfig == null) throw new AdeKnowledgeException("知识库检索配置不能为空");

        AdeRetrivedResponse response = new AdeRetrivedResponse();
        response.setPairs(new ArrayList<>());
        response.setAdes(new ArrayList<>());

        List<String> references = new ArrayList<>();  //切割后的句子
        //STEP 1: 判断EMR分割类型
        IAbstractChunkRequest chunkRequest = null;
        if (IAbstractChunkRequest.SPLIT_PATTERN.equals(retriveConfig.getChunkType())) {
            chunkRequest = new ChunkByPatternRequest();
            IChunkService chunkService = new ChunkByPatternService();
            List<String> sentences = chunkService.chunk(emr, chunkRequest);
            references.addAll(sentences);
        } else if (IAbstractChunkRequest.SPLIT_MARK.equals(retriveConfig.getChunkType())) {
            chunkRequest = ChunkByMarkRequest.getDefault();
            IChunkService chunkService = new ChunkByMarkService();
            List<String> sentences = chunkService.chunk(emr, chunkRequest);
            references.addAll(sentences);
        } else if (IAbstractChunkRequest.SPLIT_NONE.equals(retriveConfig.getChunkType())) {
            references.add(emr);
        } else {
            throw new AdeKnowledgeException("未知的分割类型:" + retriveConfig.getChunkType());
        }

        if(retriveConfig.isDebug()) {
            //step2: 每个chunk分析
            for (String reference : references) {
                AdeRetrivedPair pair = new AdeRetrivedPair();
                pair.setChunk(reference);
                pair.setAdes(new ArrayList<>());
                EmbeddingRequest embeddingRequest = new EmbeddingRequest(retriveConfig.getEmbedModel(), Arrays.asList(new String[]{reference}));
                log.info("retriveConfig:{}", JSONObject.toJSONString(retriveConfig));
                log.info("embeddingRequest:{}", JSONObject.toJSONString(embeddingRequest));
                EmbeddingResponse embeddingResponse = ollamaAgent.embed(embeddingRequest);
                Float[][] vector = embeddingResponse.getEmbeddings();
                //从向量库批量检索
                List<AdeMilvusResponse> responses = milvusAgent.search(retriveConfig.getCollection(), vector, retriveConfig.getSimScoreThreshold(), retriveConfig.getTopN());
                for (AdeMilvusResponse resp : responses) {
                    AdeRetrivedKnowledge r = AdeRetrivedKnowledge.fromMilvusObject(resp);
                    pair.getAdes().add(r);
                    response.getAdes().add(r);
                }
                Collections.sort(pair.getAdes()); //排序 从大到小按分数排序
                response.getPairs().add(pair);
            }
        }else{
            EmbeddingRequest embeddingRequest1 = new EmbeddingRequest(retriveConfig.getEmbedModel(), references);
            EmbeddingResponse embeddingResponse1 = ollamaAgent.embed(embeddingRequest1);
            Float[][] vector1 = embeddingResponse1.getEmbeddings();
            List<AdeMilvusResponse> responses1 = milvusAgent.search(retriveConfig.getCollection(), vector1, retriveConfig.getSimScoreThreshold(), retriveConfig.getTopN());
            for (AdeMilvusResponse resp : responses1) {
                AdeRetrivedKnowledge r = AdeRetrivedKnowledge.fromMilvusObject(resp);
                response.getAdes().add(r);
            }
        }

//        //Step2: 基于向量数据库 检索相似性ADE知识
//        //将查询的文本转换成向量
//        EmbeddingRequest embeddingRequest = new EmbeddingRequest(retriveConfig.getEmbedModel(), references);  //批量请求向量化
//        EmbeddingResponse embeddingResponse = ollamaAgent.embed(embeddingRequest);
//        Float[][] vector = embeddingResponse.getEmbeddings();
//        //从向量库批量检索
//        List<AdeMilvusResponse> responses = milvusAgent.search(retriveConfig.getCollection(), vector, retriveConfig.getSimScoreThreshold(), retriveConfig.getTopN());
//        Collections.sort(responses); //从打到小排序
//        int max = retriveConfig.getTopN();  //假设检索top 10 ,可能会有100个出来， 因为有多个语句检索每个语句都会返回top10,  我们最大值保留两倍的topn
//        if (responses.size() > max) responses = responses.subList(0, max);
//
//        //Step3: 组装RetrievedAcknowledge回去
//        List<AdeRetrivedKnowledge> raks = new ArrayList<>();
//        Map<String, AdeRetrivedKnowledge> map = new HashMap<>();
//        for (AdeMilvusResponse resp : responses) {
//            AdeRetrivedKnowledge r = AdeRetrivedKnowledge.fromMilvusObject(resp);
//            map.put(r.getKnowledge().getText(), r);
//            raks.add(r);
//        }
//
//        //Step4: 判断是否需要Rerank   必须确保有rerank对象 即responses不能为空
//        if (rerankConfig != null && responses.size() > 0) {
//            raks.clear();
//            List<String> docs = new ArrayList<>();
//            for (AdeMilvusResponse resp : responses) docs.add(resp.getText());
//            int min = docs.size();
//            if (rerankConfig.getTopN() < min)
//                min = rerankConfig.getTopN();  //相关性检索结果返回的top N    如果候选项有10个 topn=5, 那么topn就是5， 如果候选项有3个，topn是5 那么topn就是3 。 取两个值的最小值
//            RerankRequest rerankRequest = new RerankRequest(rerankConfig.getRerankModel(), emr, docs, min);
//            RerankResponse rerankResponse = qianfanAgent.reran(rerankRequest);
////            System.out.println(rerankResponse+"@@@@@@@@@@@@@@@@@");
//            if (rerankResponse != null)
//                for (RerankResponse.RerankDocument doc : rerankResponse.getResults()) {
//                    AdeRetrivedKnowledge tmp = map.get(doc.getDocument());
//                    tmp.setReScore(doc.getRelevance_score());
//                    if (tmp.getReScore() >= rerankConfig.getReScoreThreshold())  //满足相关性得分的才要
//                        raks.add(tmp);
//                }
//        }
//        return raks;
        response.setAdes(filterRepartAndSort(response.getAdes(), retriveConfig.getTopN()));
        return response;
    }


    private List<AdeRetrivedKnowledge> filterRepartAndSort(List<AdeRetrivedKnowledge> ades, int topN){
        List<AdeRetrivedKnowledge> results = new ArrayList<>();
        Map<String,Byte> flag = new HashMap<>();
        for (AdeRetrivedKnowledge ade : ades) {
            if(flag.get(ade.getKnowledge().getText())==null){
                    results.add(ade);
                    flag.put(ade.getKnowledge().getText(), (byte) 1);
            }
        }
        Collections.sort(results);
        if(results.size()>topN)
            results = results.subList(0, topN);
        return results;
    }
}
